package sivan.test;

import jdk.internal.org.objectweb.asm.*;
import jdk.internal.org.objectweb.asm.commons.AdviceAdapter;
import lombok.Data;
import org.springframework.cglib.core.ReflectUtils;

import java.io.FileOutputStream;

import static jdk.internal.org.objectweb.asm.Opcodes.ACC_PUBLIC;
import static jdk.internal.org.objectweb.asm.Opcodes.V1_1;

/**
 * ASM是快速操作java字节码的框架，可以读取字节码文件
 * 和动态生成字节码文件
 */
public class ASMDemo {

    public static void main(String[] args) {
        try {
            ClassReader reader = new ClassReader(HelloWorld.class.getName());

            ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
            writer.visit(V1_1, ACC_PUBLIC, "ASMDemo$HelloWorld", null, "java/lang/Object", null);
            // 3. 创建自定义的 ClassVisitor 对象
            ClassVisitor change = new ChangeVisitor(writer);

            reader.accept(change, ClassReader.EXPAND_FRAMES);

            byte[] code = writer.toByteArray();
            String path = ASMDemo.class.getProtectionDomain().getCodeSource().getLocation().getPath();
            FileOutputStream fos = new FileOutputStream(path + "/sivan/test/ASMDemo$HelloWorld.class");
            fos.write(code);
            fos.close();

            Class helloProxy = ReflectUtils.defineClass("sivan.test.HelloWorld", code, new MyClassLoader<>());
            Object o = helloProxy.newInstance();
            helloProxy.getMethod("sayHello").invoke(o, null);

            Class<?> data = Class.forName(HelloWorld.class.getName());
            data.getMethod("sayHello").invoke(data.newInstance(), null);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static class MyClassLoader<T> extends ClassLoader {

        @SuppressWarnings("unchecked")
        public Class<T> defineMyClass(String name, byte[] bs, int off, int len) {
            return (Class<T>) super.defineClass(name, bs, off, len);
        }
    }

    @Data
    public static class Test{
        private int id;
        private String name;
    }

    static class ChangeVisitor extends ClassVisitor {

        ChangeVisitor(ClassVisitor classVisitor) {
            super(Opcodes.ASM5, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
            MethodVisitor methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
            if (name.equals("<init>")) {
                return methodVisitor;
            }
            return new ChangeAdapter(Opcodes.ASM4, methodVisitor, access, name, desc);
        }
    }

    static class ChangeAdapter extends AdviceAdapter {
        private int startTimeId = -1;

        private String methodName = null;

        ChangeAdapter(int api, MethodVisitor mv, int access, String name, String desc) {
            super(api, mv, access, name, desc);
            methodName = name;
        }

        @Override
        protected void onMethodEnter() {
            super.onMethodEnter();
            startTimeId = newLocal(Type.LONG_TYPE);
            mv.visitMethodInsn(INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J", false);
            mv.visitIntInsn(LSTORE, startTimeId);
        }

        @Override
        protected void onMethodExit(int opcode) {
            super.onMethodExit(opcode);
            int durationId = newLocal(Type.LONG_TYPE);
            mv.visitMethodInsn(INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J", false);
            mv.visitVarInsn(LLOAD, startTimeId);
            mv.visitInsn(LSUB);
            mv.visitVarInsn(LSTORE, durationId);
            mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
            mv.visitTypeInsn(NEW, "java/lang/StringBuilder");
            mv.visitInsn(DUP);
            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/StringBuilder", "<init>", "()V", false);
            mv.visitLdcInsn("The cost time of " + methodName + "() is ");
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitVarInsn(LLOAD, durationId);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(J)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(" ms");
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);

        }
    }
}
